"""Module providing SSH functionality."""
import os
import paramiko
from time import sleep

from cStringIO import StringIO

# Try to issue warning on too old paramiko version
_minParamVersion = (1, 6, 1)
try:
    _paramVersion = paramiko.__version_info__
except AttributeError:
    _paramVersion = None

if not _paramVersion or _paramVersion < _minParamVersion:
    _paramVersion = ".".join(_paramVersion) if _paramVersion else "UNKNOWN"
    raise Warning("Probably unsupported paramiko version: %s" % _paramVersion)

class _SshPopen(object):
    """Object that is thought to emulate subprocess.Popen on top
    of SSH channel.

    """

    returncode = None
    _command = None
    stdout = stderr = None
    # note that tuples are immutable!
    _stdoutHandlers = _stderrHandlers = ()
    _waitCallbacks = _pollCallbacks = ()
    _deathCallbacks = ()
    alive = False

    def __init__(self, channel, command, autoStart=True,
        useDefaultStdout=True, useDefaultStderr=True
        ):
        self._chan = channel
        self._command = command
        self.stdin = self._chan.makefile("wb")
        self.stdout = StringIO()
        self.stderr = StringIO()

        if useDefaultStdout:
            self.addStdoutHandler(
                lambda data: self._appendToStream(self.stdout, data))
        if useDefaultStderr:
            self.addStdoutHandler(
                lambda data: self._appendToStream(self.stderr, data))
        if autoStart:
            self.start()

    def start(self):
        self._chan.exec_command(self._command)
        self.alive = True

    def addStdoutHandler(self, handler):
        self._stdoutHandlers += (handler, )

    def addStderrHandler(self, handler):
        self._stderrHandlers += (handler, )

    def getCommand(self):
        return self._command

    def __repr__(self):
        return "<%s:%r>" % (self.__class__.__name__, self._command)

    @staticmethod
    def _appendToStream(stream, data):
        _lastPos = stream.tell()
        stream.seek(0, os.SEEK_END)
        stream.write(data)
        stream.seek(_lastPos, os.SEEK_SET)

    def _recvData(self):
        """Read available data from channel to stdon/stderr.

        Note that it is only method that writes data to given streams,
        therefore it must be called from time to time.

        """
        _chan = self._chan
        _data = ""
        while _chan.recv_ready():
            _data += _chan.recv(1024)
        if _data:
            [_handle(_data) for _handle in self._stdoutHandlers]
        _data = ""
        while _chan.recv_stderr_ready():
            _data += _chan.recv_stderr(1024)
        if _data:
            [_handle(_data) for _handle in self._stderrHandlers]

    def _exitStatusReady(self):
        return self._chan.closed or self._chan.status_event.isSet()

    def _updateRetcode(self):
        if self.returncode is None and self._exitStatusReady():
            self.returncode = self._chan.recv_exit_status()

    def addDeathCallback(self, cb):
        self._deathCallbacks += (cb, )

    def poll(self):
        self._updateRetcode()
        self._recvData()
        [_cb() for _cb in self._pollCallbacks]
        if self.alive and self.returncode is not None:
            [_cb() for _cb in self._deathCallbacks]
            self.alive = False
        return self.returncode

    def addPollCallback(self, cb):
        self._pollCallbacks += (cb, )

    def addWaitCallback(self, cb):
        self._waitCallbacks += (cb, )

    def wait(self):
        """Block untill process finishes.

        Call given callback from time to time while blocked.
        """
        [_cb() for _cb in self._waitCallbacks]
        while self.poll() is None:
            sleep(0.1)
        return self.poll()

    def shutdown(self):
        if self.poll() is None:
            self._chan.shutdown(0)

class SshConnection(object):

    _host = _port = _user = _pass = None
    _transport = _channel = _pty = None
    _shellPrefix = None

    def __init__(self, host, port, user, password):
        self._host = host
        self._port = port
        self._user = user
        self._pass = password

    def connect(self):
        """Perform a connection."""
        _transport = paramiko.Transport((self._host, self._port))
        _transport.use_compression(True)
        # SSH on the Freerunner is quite slow, so
        # paramico could timeout sometimes when connection
        # is initiated.
        #
        # Changing timeout value to 120 seconds (two minutes)
        # is likely to relax (but not solve) this problem.
        _transport.banner_timeout = 120
        _transport.connect(username=self._user, password=self._pass)
        self._transport = _transport

    def popen(self, command, **kwargs):
        """Return object mimicking non-blocking subprocess.Popen."""
        _chan = self._transport.open_session()
        return _SshPopen(_chan, command, **kwargs)

    def close(self):
        """Terminate the connection."""
        self._transport.close()

# vim: set sts=4 sw=4 et :
